{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2016 Google LLC  \n",
    "  \n",
    " Licensed under the Apache License, Version 2.0 (the \"License\");  \n",
    " you may not use this file except in compliance with the License.  \n",
    " You may obtain a copy of the License at  \n",
    "  \n",
    "     http://www.apache.org/licenses/LICENSE-2.0  \n",
    "  \n",
    " Unless required by applicable law or agreed to in writing, software  \n",
    " distributed under the License is distributed on an \"AS IS\" BASIS,  \n",
    " WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  \n",
    " See the License for the specific language governing permissions and  \n",
    " limitations under the License."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the dependency .py files, if any.\n",
    "! git clone https://github.com/GoogleCloudPlatform/cloudml-samples.git\n",
    "! cp cloudml-samples/census/customestimator/trainer/* .\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os\n",
    "import tensorflow as tf\n",
    "import model as model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_and_evaluate(args):\n",
    "  \"\"\"Run the training and evaluate using the high level API.\"\"\"\n",
    "  train_input = lambda: model.input_fn(\n",
    "      args.train_files,\n",
    "      num_epochs=args.num_epochs,\n",
    "      batch_size=args.train_batch_size\n",
    "  )\n",
    "\n",
    "  # Don't shuffle evaluation data.\n",
    "  eval_input = lambda: model.input_fn(\n",
    "    args.eval_files,\n",
    "      batch_size=args.eval_batch_size,\n",
    "      shuffle=False\n",
    "  )\n",
    "\n",
    "  train_spec = tf.estimator.TrainSpec(\n",
    "      train_input, max_steps=args.train_steps)\n",
    "\n",
    "  exporter = tf.estimator.FinalExporter(\n",
    "      'census', model.SERVING_FUNCTIONS[args.export_format])\n",
    "  eval_spec = tf.estimator.EvalSpec(\n",
    "      eval_input,\n",
    "      steps=args.eval_steps,\n",
    "      exporters=[exporter],\n",
    "      name='census-eval')\n",
    "\n",
    "  model_fn = model.generate_model_fn(\n",
    "      embedding_size=args.embedding_size,\n",
    "      # Construct layers sizes with exponential decay.\n",
    "      hidden_units=[\n",
    "          max(2, int(args.first_layer_size * args.scale_factor**i))\n",
    "          for i in range(args.num_layers)\n",
    "      ],\n",
    "      learning_rate=args.learning_rate)\n",
    "\n",
    "  estimator = tf.estimator.Estimator(\n",
    "      model_fn=model_fn, model_dir=args.job_dir)\n",
    "  tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "# Input Arguments.\n",
    "parser.add_argument(\n",
    "    '--train-files',\n",
    "    help='GCS file or local paths to training data',\n",
    "    nargs='+',\n",
    "    default='gs://cloud-samples-data/ml-engine/census/data/adult.data.csv')\n",
    "parser.add_argument(\n",
    "    '--eval-files',\n",
    "    help='GCS file or local paths to evaluation data',\n",
    "    nargs='+',\n",
    "    default='gs://cloud-samples-data/ml-engine/census/data/adult.test.csv')\n",
    "parser.add_argument(\n",
    "    '--job-dir',\n",
    "    help='GCS location to write checkpoints and export models',\n",
    "    default='/tmp/census-customestimator')\n",
    "parser.add_argument(\n",
    "    '--num-epochs',\n",
    "    help=\"\"\"\\\n",
    "    Maximum number of training data epochs on which to train.\n",
    "    If both --max-steps and --num-epochs are specified,\n",
    "    the training job will run for --max-steps or --num-epochs,\n",
    "    whichever occurs first. If unspecified will run for --max-steps.\\\n",
    "    \"\"\",\n",
    "    type=int,\n",
    ")\n",
    "parser.add_argument(\n",
    "    '--train-batch-size',\n",
    "    help='Batch size for training steps',\n",
    "    type=int,\n",
    "    default=40)\n",
    "parser.add_argument(\n",
    "    '--eval-batch-size',\n",
    "    help='Batch size for evaluation steps',\n",
    "    type=int,\n",
    "    default=40)\n",
    "parser.add_argument(\n",
    "    '--embedding-size',\n",
    "    help='Number of embedding dimensions for categorical columns',\n",
    "    default=8,\n",
    "    type=int)\n",
    "parser.add_argument(\n",
    "    '--learning-rate',\n",
    "    help='Learning rate for the optimizer',\n",
    "    default=0.1,\n",
    "    type=float)\n",
    "parser.add_argument(\n",
    "    '--first-layer-size',\n",
    "    help='Number of nodes in the first layer of the DNN',\n",
    "    default=100,\n",
    "    type=int)\n",
    "parser.add_argument(\n",
    "    '--num-layers',\n",
    "    help='Number of layers in the DNN',\n",
    "    default=4, type=int)\n",
    "parser.add_argument(\n",
    "    '--scale-factor',\n",
    "    help='How quickly should the size of the layers in the DNN decay',\n",
    "    default=0.7,\n",
    "    type=float)\n",
    "parser.add_argument(\n",
    "    '--train-steps',\n",
    "    help=\"\"\"\\\n",
    "    Steps to run the training job for. If --num-epochs is not specified,\n",
    "    this must be. Otherwise the training job will run indefinitely.\\\n",
    "    \"\"\",\n",
    "    default=100,\n",
    "    type=int)\n",
    "parser.add_argument(\n",
    "    '--eval-steps',\n",
    "    help=\"\"\"\\\n",
    "    Number of steps to run evalution for at each checkpoint.\n",
    "    If unspecified will run until the input from --eval-files is exhausted\"\"\",\n",
    "    default=None,\n",
    "    type=int)\n",
    "parser.add_argument(\n",
    "    '--export-format',\n",
    "    help='The input format of the exported SavedModel binary',\n",
    "    choices=['JSON', 'CSV', 'EXAMPLE'],\n",
    "    default='JSON')\n",
    "parser.add_argument(\n",
    "    '--verbosity',\n",
    "    choices=['DEBUG', 'ERROR', 'FATAL', 'INFO', 'WARN'],\n",
    "    default='INFO',\n",
    "    help='Set logging verbosity')\n",
    "\n",
    "args, _ = parser.parse_known_args()\n",
    "\n",
    "# Set python level verbosity.\n",
    "tf.logging.set_verbosity(args.verbosity)\n",
    "# Set C++ Graph Execution level verbosity.\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = str(\n",
    "    tf.logging.__dict__[args.verbosity] / 10)\n",
    "\n",
    "# Run the training job.\n",
    "train_and_evaluate(args)"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
